# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from experimental.prophet.graph import graph as _graph
from experimental.prophet.graph.util import to_data_type_proto


def elemwise_unary_op(func, v, node_name=None):
  v = _graph.as_variable(v)
  node = _graph.create_node('ElemwiseUnaryOp', [v],
                            name=node_name,
                            attrs={
                                'dtype': to_data_type_proto(v.dtype), 'func': func
                            })
  return node.outputs[0]


def elemwise_binary_op(func, v1, v2, node_name=None):
  v1 = _graph.as_variable(v1)
  v2 = _graph.as_variable(v2)
  node = _graph.create_node('ElemwiseBinOp', [v1, v2],
                            name=node_name,
                            attrs={
                                'dtype1': to_data_type_proto(v1.dtype),
                                'dtype2': to_data_type_proto(v2.dtype),
                                'func': func
                            })
  return node.outputs[0]


# Logical operators
def logical_not(v, node_name=None):
  return elemwise_unary_op('Not', v, node_name=node_name)


def logical_and(v1, v2, node_name=None):
  return elemwise_binary_op('And', v1, v2, node_name=node_name)


def logical_or(v1, v2, node_name=None):
  return elemwise_binary_op('Or', v1, v2, node_name=node_name)


# Comparison operators
def equal(v1, v2, node_name=None):
  return elemwise_binary_op('Equal', v1, v2, node_name=node_name)


def not_equal(v1, v2, node_name=None):
  return elemwise_binary_op('NotEqual', v1, v2, node_name=node_name)


def less(v1, v2, node_name=None):
  return elemwise_binary_op('Less', v1, v2, node_name=node_name)


def less_equal(v1, v2, node_name=None):
  return elemwise_binary_op('LessEqual', v1, v2, node_name=node_name)


def greater(v1, v2, node_name=None):
  return elemwise_binary_op('Greater', v1, v2, node_name=node_name)


def greater_equal(v1, v2, node_name=None):
  return elemwise_binary_op('GreaterEqual', v1, v2, node_name=node_name)


def isnan(v, node_name=None):
  return elemwise_unary_op('IsNan', v, node_name=node_name)


# Arithmetic operators
def neg(v, node_name=None):
  return elemwise_unary_op('Neg', v, node_name=node_name)


def add(v1, v2, node_name=None):
  return elemwise_binary_op('Add', v1, v2, node_name=node_name)


def sub(v1, v2, node_name=None):
  return elemwise_binary_op('Sub', v1, v2, node_name=node_name)


def mul(v1, v2, node_name=None):
  return elemwise_binary_op('Mul', v1, v2, node_name=node_name)


def div(v1, v2, node_name=None):
  return elemwise_binary_op('Div', v1, v2, node_name=node_name)


def mod(v1, v2, node_name=None):
  return elemwise_binary_op('Mod', v1, v2, node_name=node_name)


def fmod(v1, v2, node_name=None):
  return elemwise_binary_op('Fmod', v1, v2, node_name=node_name)


def pow(v1, v2, node_name=None):
  return elemwise_binary_op('Pow', v1, v2, node_name=node_name)


# Bitwise operators
def bitwise_not(v, node_name=None):
  return elemwise_unary_op('BitNot', v, node_name=node_name)


def bitwise_and(v1, v2, node_name=None):
  return elemwise_binary_op('BitAnd', v1, v2, node_name=node_name)


def bitwise_or(v1, v2, node_name=None):
  return elemwise_binary_op('BitOr', v1, v2, node_name=node_name)


def bitwise_xor(v1, v2, node_name=None):
  return elemwise_binary_op('BitXor', v1, v2, node_name=node_name)


def bitwise_left_shift(v1, v2, node_name=None):
  return elemwise_binary_op('BitLeftShift', v1, v2, node_name=node_name)


def bitwise_right_shift(v1, v2, node_name=None):
  return elemwise_binary_op('BitRightShift', v1, v2, node_name=node_name)


# Math functions
def abs(v, node_name=None):
  return elemwise_unary_op('Abs', v, node_name=node_name)


def exp(v, node_name=None):
  return elemwise_unary_op('Exp', v, node_name=node_name)


def exp2(v, node_name=None):
  return elemwise_unary_op('Exp2', v, node_name=node_name)


def expm1(v, node_name=None):
  return elemwise_unary_op('ExpM1', v, node_name=node_name)


def log(v, node_name=None):
  return elemwise_unary_op('Log', v, node_name=node_name)


def log2(v, node_name=None):
  return elemwise_unary_op('Log2', v, node_name=node_name)


def log10(v, node_name=None):
  return elemwise_unary_op('Log10', v, node_name=node_name)


def log1p(v, node_name=None):
  return elemwise_unary_op('Log1P', v, node_name=node_name)


def sqrt(v, node_name=None):
  return elemwise_unary_op('Sqrt', v, node_name=node_name)


def cbrt(v, node_name=None):
  return elemwise_unary_op('Cbrt', v, node_name=node_name)


def sin(v, node_name=None):
  return elemwise_unary_op('Sin', v, node_name=node_name)


def cos(v, node_name=None):
  return elemwise_unary_op('Cos', v, node_name=node_name)


def tan(v, node_name=None):
  return elemwise_unary_op('Tan', v, node_name=node_name)


def asin(v, node_name=None):
  return elemwise_unary_op('Asin', v, node_name=node_name)


def acos(v, node_name=None):
  return elemwise_unary_op('Acos', v, node_name=node_name)


def atan(v, node_name=None):
  return elemwise_unary_op('Atan', v, node_name=node_name)


def sinh(v, node_name=None):
  return elemwise_unary_op('Sinh', v, node_name=node_name)


def cosh(v, node_name=None):
  return elemwise_unary_op('Cosh', v, node_name=node_name)


def tanh(v, node_name=None):
  return elemwise_unary_op('Tanh', v, node_name=node_name)


def asinh(v, node_name=None):
  return elemwise_unary_op('Asinh', v, node_name=node_name)


def acosh(v, node_name=None):
  return elemwise_unary_op('Acosh', v, node_name=node_name)


def atanh(v, node_name=None):
  return elemwise_unary_op('Atanh', v, node_name=node_name)


def erf(v, node_name=None):
  return elemwise_unary_op('Erf', v, node_name=node_name)


def erfc(v, node_name=None):
  return elemwise_unary_op('Erfc', v, node_name=node_name)


def tgamma(v, node_name=None):
  return elemwise_unary_op('Tgamma', v, node_name=node_name)


def lgamma(v, node_name=None):
  return elemwise_unary_op('Lgamma', v, node_name=node_name)


def ceil(v, node_name=None):
  return elemwise_unary_op('Ceil', v, node_name=node_name)


def floor(v, node_name=None):
  return elemwise_unary_op('Floor', v, node_name=node_name)


def trunc(v, node_name=None):
  return elemwise_unary_op('Trunc', v, node_name=node_name)


def round(v, node_name=None):
  return elemwise_unary_op('Round', v, node_name=node_name)


def sign(v, node_name=None):
  return elemwise_unary_op('Sign', v, node_name=node_name)


def _min(v1, v2, node_name=None):
  return elemwise_binary_op('Min', v1, v2, node_name=node_name)


def min(*vars):
  assert vars
  min_v = vars[0]
  for v in vars[1:]:
    min_v = _min(min_v, v)
  return min_v


def _max(v1, v2, node_name=None):
  return elemwise_binary_op('Max', v1, v2, node_name=node_name)


def max(*vars):
  assert vars
  max_v = vars[0]
  for v in vars[1:]:
    max_v = _max(max_v, v)
  return max_v


def atan2(v1, v2, node_name=None):
  return elemwise_binary_op('Atan2', v1, v2, node_name=node_name)


# Reduction functions
def sum(v, node_name=None):
  return elemwise_unary_op('Sum', v, node_name=node_name)
